package ocr

import (
	"os"
	"path/filepath"

	"github.com/sirupsen/logrus"
)

type Config struct {
	PaddleModel
	//det_model_dir
	DetModelDir string `goblet:"det_model_dir" default:"https://paddleocr.bj.bcebos.com/PP-OCRv4/chinese/ch_PP-OCRv4_det_infer.tar"`
	//rec_model_dir
	RecModelDir string `goblet:"rec_model_dir" default:"https://paddleocr.bj.bcebos.com/PP-OCRv4/chinese/ch_PP-OCRv4_rec_infer.tar"`
	//cls_model_dir
	ClsModelDir string `goblet:"cls_model_dir" default:"https://paddleocr.bj.bcebos.com/dygraph_v2.0/ch/ch_ppocr_mobile_v2.0_cls_infer.tar"`
	//use_angle_cls
	UseAngleCls bool `goblet:"use_angle_cls"`
	//det_max_side_len
	DetMaxSideLen int `goblet:"det_max_side_len" default:"960"`
	//det_db_thresh
	DetDbThresh float64 `goblet:"det_db_thresh" default:"0.3"`
	//det_db_box_thresh
	DetDbBoxThresh float64 `goblet:"det_db_box_thresh" default:"0.5"`
	//det_db_unclip_ratio
	DetDbUnclipRatio float64 `goblet:"det_db_unclip_ratio" default:"2.0"`
	//rec_image_shape
	RecImageShape []int `goblet:"rec_image_shape" default:"[3, 32, 320]"`
	//rec_char_dict_path
	RecCharDictPath string `goblet:"rec_char_dict_path" default:"./config/ppocr_keys_v1.txt"`
	//use_space_char
	UseSpaceChar bool `goblet:"use_space_char" default:"true"`
	//rec_batch_num
	RecBatchNum int `goblet:"rec_batch_num" default:"10"`
	//max_text_length
	MaxTextLength int `goblet:"max_text_length" default:"25"`
	//rec_char_type
	RecCharType string `goblet:"rec_char_type" default:"ch"`
	//cls_image_shape
	ClsImageShape []int `goblet:"cls_image_shape" default:"[3, 48, 192]"`
	//cls_batch_num
	ClsBatchNum int `goblet:"cls_batch_num" default:"30"`
	//cls_thresh
	ClsThresh float64 `goblet:"cls_thresh" default:"0.9"`
}

func (c *Config) Check() {
	if c.ModleDir == "" {
		home, _ := os.UserHomeDir()
		c.ModleDir = filepath.Join(home, ".paddleocr")
	}
	if len(c.RecImageShape) == 0 {
		c.RecImageShape = []int{3, 32, 320}
	}
	if len(c.ClsImageShape) == 0 {
		c.ClsImageShape = []int{3, 48, 192}
	}
	logrus.WithField("model_dir", c.ModleDir).Info("use model dir")
}
